import os
import numpy as np
import scipy.optimize as opt
from tqdm import tqdm
from clean_data import dist
from functools import partial
from method import ls, phi_meta, grad_phi_meta


if __name__ == '__main__':
    anchors = [
        [0, 0, 1300],
        [5000, 0, 1700],
        [0, 5000, 1700],
        [5000, 5000, 1300]
    ]
    for C in ['正常', '异常']:
        states = [0b0000, 0b0001, 0b0010, 0b0100, 0b1000]
        for n in tqdm(range(324)):
            mat = np.loadtxt(os.path.join(f"{C}数据_clean", f"{n+1}.{C}.txt"), dtype=np.int)
            features = list()
            for t in range(len(mat)):
                b_init = ls(anchors, mat[t])
                feat = list()
                for s in states:
                    signal = [0, 0, 0, 0]
                    for j in range(4):
                        if (s >> j) & 1 == 1:
                            signal[j] = mat[t][j] - 400
                        else:
                            signal[j] = mat[t][j] + 45
                    sols = list()
                    ''' compute reference node '''
                    phi = partial(phi_meta, ref=anchors, data=signal)
                    grad_phi = partial(grad_phi_meta, ref=anchors, data=signal)
                    ref_node = opt.root(grad_phi, b_init, method='lm').x
                    feat.append(phi(ref_node)) # ref node error
                    ''' compute three test nodes '''
                    test_node_err, node_dist = list(), list()
                    for j in range(4):
                        p_anchors = [anchors[k] for k in range(4) if k != j]
                        p_data = [signal[k] for k in range(4) if k != j]
                        phi = partial(phi_meta, ref=p_anchors, data=p_data)
                        grad_phi = partial(grad_phi_meta, ref=p_anchors, data=p_data)
                        sol = opt.root(grad_phi, b_init, method='lm').x
                        sols.append(sol)
                        test_node_err.append(phi(sol)) # test node error
                        node_dist.append(dist(sol, ref_node)) # node distence
                    node_var = (np.array(sols).std(axis=0) ** 2).tolist()
                    feat = feat + test_node_err + node_dist + node_var
                features.append(np.array(feat))
            features = np.array(features)
            np.savetxt(os.path.join(f"{C}数据_clean", f"{n+1}.{C}.feat.txt"), features, fmt='%.4f', encoding='utf-8')
